package com.guava;

import com.google.common.util.concurrent.RateLimiter;
import org.springframework.http.MediaType;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.nio.charset.StandardCharsets;

/**
 * Created by fangjicai on 2020/10/28.
 */
public class RateLimiterInterceptor extends HandlerInterceptorAdapter {

    private final RateLimiter rateLimiter;

    /**
     * 通过构造函数初始化限流器
     * @param rateLimiter
     */
    public RateLimiterInterceptor(RateLimiter rateLimiter){
        super();
        this.rateLimiter=rateLimiter;
    }


    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if(this.rateLimiter.tryAcquire()){
            /**
             * 成功获取令牌
             */
            return true;
        }else{
            /**
             * 获取令牌失败
             */
            response.setCharacterEncoding(StandardCharsets.UTF_8.name());
            response.setContentType(MediaType.TEXT_PLAIN_VALUE);
            response.getWriter().write("系统繁忙!");
            return false;
        }
    }
}
